package design.leetcode_208;

public class Trie {
    private Trie[] next;
    // 记录每个单词出现的次数（题目中说明可能插入重复单词）
    private int wordCount;
    // 记录每个字符出现的次数，同时也是每个字符出现的前缀数
    private int prefixCount;

    public Trie() {
        this.next = new Trie[26];
        this.wordCount = 0;
        this.prefixCount = 0;
    }


    public void insert(String word) {
        if (word == null || word.equals("")) {
            return;
        }
        Trie node = this;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (node.next[ch - 'a'] == null) {
                node.next[ch - 'a'] = new Trie();
            }
            node = node.next[ch - 'a'];
            // 每个字符出现次数增加
            node.prefixCount++;
        }
        // 记录当前插入单词出现的次数
        node.wordCount++;
    }

    public boolean search(String word) {
        if (word == null || word.equals("")) {
            return false;
        }
        Trie node = this;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (node.next[ch - 'a'] == null) {
                return false;
            }
            node = node.next[ch - 'a'];
        }
        // 如果最后一个字符记录数 > 0 说明该单词存在于字典树中
        return node.wordCount > 0;
    }

    public int prefixNumber(String pre) {
        if (pre == null || pre.equals("")) {
            return 0;
        }
        Trie node = this;
        for (int i = 0; i < pre.length(); i++) {
            char ch = pre.charAt(i);
            if (node.next[ch - 'a'] == null) {
                return 0;
            }
            node = node.next[ch - 'a'];
        }
        return node.prefixCount;
    }

    public void delete(String word) {
        if (word == null || word.equals("") || !search(word)) {
            return;
        }
        Trie node = this;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (--node.next[ch - 'a'].prefixCount == 0) {
                node.next[ch - 'a'] = null;
                return;
            }
            node = node.next[ch - 'a'];
        }
        // 该单词出现次数减 1
        node.wordCount--;
    }
}
